#include <iostream>
#include <vector>
#include <ctime>
#include <numeric>

#include "tensor.h"
#include "covariance.h"
#include "mtrand.h"
#include "dataset.h"
#include "pdf.h"
#include <vector>
#include <sstream>
#include <fstream>

#include "TFile.h"
#include "TChain.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TProfile2D.h"
#include "TF1.h"
#include "TF2.h"
#include "TF3.h"
#include "TLorentzVector.h"
#include "TGraphErrors.h"
#include "TGraph2DErrors.h"
#include "make_list.h"

using namespace std;

//initialize random number generator
MTRand_int32 irand(time(0));

const double PI=3.14159265359;

//declare minimum pt cut
double min_pt=150.;


//Generate TGraph
//a function to filicate making TGraph for ROOT
TGraphErrors Generate_TGraph(PDF& pdf, double pt=500, string name="example")
{
  //assume that the pdf has rank 2
  if(pdf.Rank() != 2)
    {
      cout<<"ERROR: Generate_TGraph only works for template p(k, log(pt/minpt))"
	  <<endl;
      throw ;
    }

  //create a MC_Generator 
  //this is simply to dump all the values of the PDF
  PDF::MC_Generator gen=
    pdf.Generator(make_list<PDF::MC_Flag>
		  (PDF::OUTPUT, PDF::INPUT));
  
  
  //GenerateFull does NOT use random numbers,
  //it simply dumps all the values of the pdf in each bin
  vector<PDF::MC_Event> evt=
    gen.GenerateFull(make_list<double>(log(pt/min_pt)));
    
  //dump the values and store them in a TGraph
  vector<double> x,y,err;
  for(unsigned int i=0; i<evt.size(); ++i)
    {
      x.push_back(evt[i].value[0]);
      y.push_back(evt[i].weight);
      err.push_back(evt[i].err);
    }

  
  TGraphErrors result(x.size(), &x[0], &y[0], NULL, &err[0]);

  //give the plot a name
  result.SetName(name.c_str());
  result.SetTitle(name.c_str());
  
  return result;
}


int main(int argc, char **argv)
{
  if(argc < 2+1)
    {
      
      cout<<"Usage: ./madgraph_example madgraph_input output.root"
	  <<endl;

      
      cout<<"example madgraph_input: QCD4j_8TeV_MG4j"<<endl;
      cout<<"files are read as QCD4j_8TeV_MG4j?.?.root"
	  <<endl;
      cout<<"outputs are stored in output.root"<<endl;

      throw;
    }
  
  
  //extra_scale to scale up all the events
  //1e9 = scale up to pb (madgraph weights are in mb)
  double extra_scale= 1e9;
  
  //set random seed (can be custom)
  //only used for smearing templates
  /*
    DataSet::SetRandomSeed(time(0));
  */

  //set extra smoothing factor
  //i.e. scaling up effective number of statistics
  //100 means less smoothing
  //i.e. the kernel width would be as if we have 100x more data
  //0.01 means more smoothing
  //i.e. the kernel width would be as if we have 0.01x data
  
  /*
    DataSet::kernel_smoothing_factor=100;
  */

  
  //declare variables
  //for reading madgraph_input rootfiles
  const int maxsize=50;

  //number of jets
  Int_t           size;
  
  //weight of event
  Double_t         w;   //[size]
  
  //kinematic + substructure variables for jets
  Float_t         pt[maxsize];   //[size]
  Float_t         eta[maxsize];
  Float_t         phi[maxsize];
  Float_t         m[maxsize];   //[size]
  Float_t         t1[maxsize];
  Float_t         t2[maxsize];
  Float_t         t3[maxsize];


  //declare DataSet for creating PDFs
  //tau21 and pt
  DataSet dataset_tau21_pt(make_list<double>(0, 0),
			   make_list<double>(1.0, 2.8),
			   make_list<double>(0.001, 0.002));


  //now process madgraph events
  char alist[3]={'A', 'B', 'C'};

  //here we proceed to the training
  for(unsigned int k=1; k<=12; k++)
    for(unsigned int ka=0; ka<3; ka++)
      {

	//Declare a TChain for reading rootfiles
	TChain chain_madgraph("Ntuple");

	//get the actual file name
	stringstream sout;
	sout<<argv[1]<<alist[ka]<<"."<<k<<".root";

	//load the file to TChain
	//and also load the TBranches
	chain_madgraph.Add(sout.str().c_str());
	chain_madgraph.SetBranchAddress("size", &size);
	chain_madgraph.SetBranchAddress("w", &w);
	chain_madgraph.SetBranchAddress("pt", pt);
	chain_madgraph.SetBranchAddress("eta", eta);
	chain_madgraph.SetBranchAddress("phi", phi);
	chain_madgraph.SetBranchAddress("m", m);
	chain_madgraph.SetBranchAddress("tau1", t1);
	chain_madgraph.SetBranchAddress("tau2", t2);
	chain_madgraph.SetBranchAddress("tau3", t3);
	
	//rescale the weights by the number of events in each bin
	double rescale=1.0/chain_madgraph.GetEntries();
	
	//now we are ready to loop over all events in this file
	for(unsigned int i=0; i<chain_madgraph.GetEntries(); ++i)
	  {

	    //basic progress bar
	    if(i%10000==0)
	      cout<<"\r\033[K"<<floor((i/(chain_madgraph.GetEntries()*12.*3)  +
				       (3*(k-1)+ka)/(12.*3)
				       )*10000)*0.01
		  <<"% trained..."
		  <<flush;
	  
	    //read the file at event i
	    chain_madgraph.GetEntry(i);

	    //require at least two jets
	    if(size<2) continue;

	    //rescale the weight
	    w*=rescale;
	    
	    int n_real=0;

	    //loop over all jets
	    //only look at first two jets for now
	    for(int n=0; n<size && n_real<2; ++n)
	      {
		//perform basic cut on the jets
		//tau1, tau2 must be non-zero
		//pt > min_pt
		//mass > 20.

		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;
		if(pt[n] < min_pt) continue;
		if(m[n]<20) continue;

		//reparameterize the variables
		double logpt=log(pt[n]/min_pt);
		double t21=t2[n]/t1[n];

		//mass / pt template not used in example
		/*
		  double mpt=-log10(m[n]/pt[n]);
		*/

		//populate the dataset
		dataset_tau21_pt.Fill(make_list<double>(t21,logpt), w);

	      }//done looping through all (first two) jets 
	
	  }//done looping through all events in a single root file
	
      }//done loop through all root files (All the bins)
  
  //update progress bar
  cout<<"\r\033[K100% trained."<<endl<<flush;


  //turn on bias only in the error computation
  //otherwise the error includes bias^2 + variance for the templates
  /*
    DataSet::b_bias_only = true;
  */

  //example on getting the values of PDF
  //PDF will be normalized to the domain specified
  PDF PDF_tau21_pt=dataset_tau21_pt.ComputePDF();

  //we can evaluate the pdf at different locations
  //if out of bounds, the end point values will be returned
  cout<<"Template at pt=500 GeV and tau21=0.3 : "
      <<PDF_tau21_pt(make_list<double>(0.3, log(500/min_pt)))<<endl;

  //one can compute the error too
  double err=0;
  cout<<"Template at pt=1000 GeV and tau21=0.5 : "
      <<PDF_tau21_pt(make_list<double>(0.5, log(1000/min_pt)), &err)<<endl;
  cout<<"error: "<<err<<endl;

  //output results
  TFile fileout(argv[2], "RECREATE");

  //I've created a function that automatically creates dumps the values
  //of the template and store it in a graph

  TGraphErrors graph1=Generate_TGraph(PDF_tau21_pt, 500, "tau21_pt500");
  graph1.Write();

  TGraphErrors graph2=Generate_TGraph(PDF_tau21_pt, 300, "tau21_pt300");
  graph2.Write();

  //close the root file
  fileout.Close();

  cout<<"Program ends."<<endl;
  return 1;
}
  
